package com.small.xx.ai.spring.chat.controller;

import com.small.xx.ai.spring.toolcall.component.time.method.TimeTools;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.http.MediaType;
import org.springframework.stereotype.Controller;
import org.springframework.ui.Model;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import reactor.core.publisher.Flux;

import java.io.IOException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

@Controller
public class AIChatController {

    private final ChatClient chatClient;
    private final ExecutorService executorService = Executors.newSingleThreadExecutor();

    public AIChatController(ChatClient.Builder builder) {
        this.chatClient = builder.build();
    }

    @GetMapping("/ai")
    public String chatHome(Model model) {
        model.addAttribute("title", "AI智能助手");
        model.addAttribute("welcomeMessage", "欢迎使用AI智能助手！请输入您的问题，我将尽力为您解答。");
        return "ai/chat";
    }

    @GetMapping("/ai/chat")
    public SseEmitter chat(@RequestParam(value = "message", defaultValue = "") String message) {
        SseEmitter emitter = new SseEmitter(Long.MAX_VALUE);
        
        if (message.isEmpty()) {
            try {
                emitter.send(SseEmitter.event().name("end").data(""));
                emitter.complete();
            } catch (IOException e) {
                emitter.completeWithError(e);
            }
            return emitter;
        }

        executorService.execute(() -> {
            try {
                Flux<String> stream = chatClient.prompt(message).stream().content();
                
                stream.subscribe(
                    data -> {
                        try {
                            // 发送数据块
                            emitter.send(SseEmitter.event().name("message").data(data));
                        } catch (IOException e) {
                            emitter.completeWithError(e);
                        }
                    },
                    error -> {
                        emitter.completeWithError(error);
                    },
                    () -> {
                        try {
                            // 发送结束信号
                            emitter.send(SseEmitter.event().name("end").data(""));
                            emitter.complete();
                        } catch (IOException e) {
                            emitter.completeWithError(e);
                        }
                    }
                );
            } catch (Exception e) {
                emitter.completeWithError(e);
            }
        });

        return emitter;
    }
    
    // 添加一个普通响应的方法用于预览
    @PostMapping("/ai/chat/preview")
    public String chatPreview(@RequestParam(value = "message", defaultValue = "") String message,
                              Model model) {
        if (!message.isEmpty()) {
            String response = chatClient.prompt(message).call().content();
            model.addAttribute("response", response);
        }
        
        model.addAttribute("message", message);
        model.addAttribute("title", "AI智能助手");
        return "ai/chat :: preview"; // 返回片段
    }
}